import argparse
import os
from numpy.lib.type_check import real
os.environ['CUDA_VISIBLE_DEVICES'] = "1"
import numpy as np
import time
import datetime
import sys
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader
from torchvision import datasets
from torch.autograd import Variable
# 自定义模块
from models.model import *
# from models.CasNet import *
# from models.ResNet import *
from models.ResNetUnet import *
# from models.Unet import *
from datasets import *
from models.pix2pix import *
# from models.feature import *
from models.transGAN import TransGenerator
import torch
from torch.optim.lr_scheduler import LambdaLR, StepLR
import piq
from models.atten_unet import AttU_Net, R2AttU_Net

from tensorboardX import SummaryWriter

writer = SummaryWriter()


path = "/mnt/4t/ljt/datasets/madic"


parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=2001, help="number of epochs of training")
parser.add_argument("--step", type=int, default=200, help="number of step of training and validation")

parser.add_argument("--dataset_name", type=str, default="atten_unet", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=4, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=0, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height", type=int, default=256, help="size of image height")
parser.add_argument("--img_width", type=int, default=256, help="size of image width")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument(
    "--sample_interval", type=int, default=500, help="interval between sampling of images from generators"
)
parser.add_argument("--checkpoint_interval", type=int, default=100, help="interval between model checkpoints")
parser.add_argument("--lambda_pixel", type=int, default=100, help="the value of labmda pixel")
parser.add_argument("--use_PixelGAN", type=bool, default=False, help="the value of labmda pixel")
opt = parser.parse_args()
# print(opt)
os.makedirs("images/%s" % opt.dataset_name, exist_ok=True)
os.makedirs("saved_models/%s" % opt.dataset_name, exist_ok=True)

# # 释放闲置内存
# if hasattr(torch.cuda, 'empty_cache'):
# 	torch.cuda.empty_cache()

cuda = True if torch.cuda.is_available() else False
# cuda = False

# os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
# os.environ['CUDA_VISIBLE_DEVICES'] = '1'
# cuda = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 损失函数

criterion_GAN = torch.nn.MSELoss()
'''L1Loss: 平均绝对误差（MAE）'''
# criterion_pixelwise = torch.nn.L1Loss()
criterion_pixelwise = torch.nn.SmoothL1Loss()

# Loss weight of L1 pixel-wise loss between translated image and real image
# 翻译图像和真实图像
lambda_pixel = opt.lambda_pixel

# Calculate output of image discriminator (PatchGAN)

patch = (1, opt.img_height // 2 ** 4, opt.img_width // 2 ** 4)
# if opt.use_PixelGAN:
#     patch = (1, 256, 256)

# Initialize generator and discriminator
# generator = ResnetGenerator(3, 3)
# generator = UnetGenerator(3, 3, 8) # 采用这个生成器，生成的图像更有“想象力”，但很糊、

# generator = UnetGenerator(1, 1, 6)
# generator = UNET()
# discriminator = Discriminator()

# generator = ResNet(Bottleneck, [3, 4, 6, 3, 3, 6, 4, 3])
# generator = CasNet()
# generator = UNET(in_channels=1, out_channels=1)
# generator = TransGenerator()
generator = R2AttU_Net()
discriminator = Discriminator()
# discriminator = TaoDiscriminator(2, 1, 8)
# discriminator = PixelDiscriminator()
# generator = GeneratorUNet()
# if cuda:
#     generator = generator.cuda()
#     discriminator = discriminator.cuda()
#     # discriminator = nn.DataParallel(discriminator)
#     criterion_GAN.cuda()
#     criterion_pixelwise.cuda()

if cuda:
    generator = generator.cuda()
    # generator = nn.DataParallel(generator)
    discriminator = discriminator.cuda()
    # discriminator = nn.DataParallel(discriminator)
    criterion_GAN.cuda()
    # criterion_GAN = nn.DataParallel(criterion_GAN)
    criterion_pixelwise.cuda()
    # criterion_pixelwise = nn.DataParallel(criterion_pixelwise)

if opt.epoch != 0:
    # 加载已经训练好的模型
    generator.load_state_dict(torch.load("saved_models/%s/generator_%d.pth" % (opt.dataset_name, opt.epoch)))
    discriminator.load_state_dict(torch.load("saved_models/%s/discriminator_%d.pth" % (opt.dataset_name, opt.epoch)))
else:
    # 初始化权重
    generator.apply(weights_init_normal)
    discriminator.apply(weights_init_normal)

# 优化器
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))

# 设置数据加载器
# transforms_ = [
#     transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),
#     transforms.ToTensor(),
#     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
# ]

transforms_ = [
    transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize(0.5, 0.5),
]

dataloader = DataLoader(
    ImageDataset(path, transforms_=transforms_),
    batch_size=opt.batch_size,
    shuffle=True,
    num_workers=opt.n_cpu,
    drop_last=True,
)

val_dataloader = DataLoader(
    ImageDataset(path, transforms_=transforms_, mode="val"),
    batch_size=opt.batch_size,
    shuffle=True,
    num_workers=opt.n_cpu,
    drop_last=True,
)

# Tensor type
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor


# ----------------------
#  图片取样
# ----------------------
def sample_images(batches_done):
    """Saves a generated sample from the validation set"""
    imgs = next(iter(val_dataloader))
    real_A = Variable(imgs["A"].type(Tensor))
    real_B = Variable(imgs["B"].type(Tensor))
    fake_B = generator(real_A)
    # mse, psnr, ssim = eval(real_B, fake_B)
    # print(mse, psnr, ssim)
    # plt.imshow(real_A)
    # plt.show()
    img_sample = torch.cat((real_A.data, fake_B.data, real_B.data), -2)
    save_image(img_sample, "images/%s/%s.png" % (opt.dataset_name, batches_done), nrow=4, normalize=True)



"""
将传入的tensor归一化到[0,1]
"""
def norm_Zero2One(x, stratic=False):
    if torch.min(x) >= 0:
        norm_x = x / (torch.max(x) - torch.min(x))
    elif torch.min(x) < 0:
        norm_x = (x - torch.min(x)) / (torch.max(x) - torch.min(x))
    # 在训练的时候发现有时候会出现越界的情况，比如x.max() = 1.0002
    # 在linux系统上未发生，在win10系统上出现过。
    if stratic:
        norm_x[norm_x > 1] = 1
        norm_x[norm_x < 0] = 0
    return norm_x


# ----------------------
#  验证集表现
# ----------------------
def val(epoch):
    # generator.eval()
    val_all_ms_ssim = 0.0
    val_all_ssim = 0.0
    val_all_psnr = 0.0
    val_all_vif = 0.0
    val_all_mse = 0.0
    val_dataiter = iter(val_dataloader)
    # for val_i, val_batch in enumerate(val_dataloader):
    for val_i in range(opt.step):
        val_batch = val_dataiter.next()
        val_real_A = Variable(val_batch["A"].type(Tensor), requires_grad=False)
        val_real_B = Variable(val_batch["B"].type(Tensor), requires_grad=False)
        # with torch.no_grad():
        val_fake_B = generator(val_real_A)
        val_tmp_real_B = norm_Zero2One(val_real_B, stratic=True)
        val_tmp_fake_B = norm_Zero2One(val_fake_B, stratic=True)

        val_ms_ssim = piq.multi_scale_ssim(val_tmp_fake_B, val_tmp_real_B, data_range=1.)
        val_ssim = piq.ssim(val_tmp_fake_B, val_tmp_real_B, data_range=1.)
        val_psnr = piq.psnr(val_tmp_fake_B, val_tmp_real_B, data_range=1., reduction='none')
        val_vif = piq.vif_p(val_tmp_real_B, val_tmp_fake_B, data_range=1.)
        val_mse = (torch.abs(val_tmp_fake_B - val_tmp_real_B) ** 2)

        val_all_ms_ssim += val_ms_ssim.item()
        val_all_ssim += val_ssim.item()
        val_all_psnr += val_psnr.mean().item()
        val_all_vif += val_vif.item()
        val_all_mse += val_mse.mean()

        # torch.cuda.empty_cache()

    val_num = val_i + 1
    # print("\n val_num: {}".format(val_num))
    # print('val_ssim', val_all_ssim / val_num, epoch)
    # print('val_psnr', val_all_psnr / val_num, epoch)
    # print('val_vif', val_all_vif / val_num, epoch)
    # print('val_mse', val_all_mse / val_num, epoch)

    writer.add_scalar('val_ssim', val_all_ssim / val_num, epoch)
    writer.add_scalar('val_ms_ssim', val_all_ms_ssim / val_num, epoch)
    writer.add_scalar('val_psnr', val_all_psnr / val_num, epoch)
    writer.add_scalar('val_vif', val_all_vif / val_num, epoch)
    writer.add_scalar('val_mse', val_all_mse / val_num, epoch)

# exp_log = open('exp_log/test_628.txt', 'a', encoding='utf-8')
# if epoch == 0:
#     exp_log.write(time.asctime(time.localtime(time.time())))
#     exp_log.write("\n超参: " + str(opt))
# exp_log.write("\n Epoch: {} \t val_mse: {} \t val_psnr: {}".format(epoch, val_mse, val_psnr))
# print("写入：val_mse: {}, val_psnr: {} ".format(val_mse, val_psnr))


# ----------------------
#  学习率改变函数
# ----------------------

def adjust_learning_rate(epoch):
    lr = opt.lr

    # if epoch > 610:
    #     lr = lr / 10
    # elif epoch > 510:
    #     lr = lr / 10
    # elif epoch > 410:
    #     lr = lr / 10
    # elif epoch > 310:
    #     lr = lr / 10
    # elif epoch > 210:
    #     lr = lr / 10
    # if epoch > 110:
    #     lr = lr / 10
    # 
    decay =  (lr / 2000) * epoch
    lr = lr - decay
    if lr <= 0:
        lr = 0
    
    writer.add_scalar('learning_rate', lr, epoch)

    for G_param_group in optimizer_G.param_groups:
        G_param_group["lr"] = lr
    for D_param_group in optimizer_D.param_groups:
        D_param_group["lr"] = lr

epoch = opt.n_epochs

# lambda1 = lambda epoch: epoch // 30
lambda1 = lambda epoch: 0.998 ** epoch

# scheduler_1 = StepLR(optimizer_1, step_size=3, gamma=0.1)
scheduler_G = LambdaLR(optimizer_G, lr_lambda=lambda1)
scheduler_D = LambdaLR(optimizer_D, lr_lambda=lambda1)



# exp_log = open('exp_log/opt.txt', 'a', encoding='utf-8')
# exp_log.write(time.asctime(time.localtime(time.time())))
# exp_log.write("\n超参: " + str(opt) + "\n")

# ----------
#  Training
# ----------
if __name__ == '__main__':
    prev_time = time.time()
    for epoch in range(opt.epoch, opt.n_epochs):
        all_loss_G = 0.0
        all_loss_D = 0.0
        all_loss_GAN = 0.0
        all_loss_pixel = 0.0
        all_loss_vif = 0.0
        all_ssim = 0.0
        all_ms_ssim = 0.0
        all_psnr = 0.0
        all_vif = 0.0
        all_mse = 0.0
        all_edge_loss = []
        dataiter = dataloader.__iter__()
        #for i, batch in enumerate(dataloader):
        for i in range(opt.step):
            batch = dataiter.__next__()
            # Model inputs
            real_A = Variable(batch["A"].type(Tensor))
            real_B = Variable(batch["B"].type(Tensor))
            # Adversarial ground truths
            valid = Variable(Tensor(np.ones((real_A.size(0), *patch))), requires_grad=False)  # [1, 1, 16, 16]
            fake = Variable(Tensor(np.zeros((real_A.size(0), *patch))), requires_grad=False)
            # valid = Variable(Tensor(np.ones((155, *patch))), requires_grad=False)
            # fake = Variable(Tensor(np.zeros((155, *patch))), requires_grad=False)

            # ------------------
            #  训练生成器
            # ------------------

            optimizer_G.zero_grad()
    

            # GAN loss
            # L1 loss
            # fake_B_1 = 0.8 * generator(real_A)
            # fake_B_2 = 0.1 * generator(fake_B_1)
            # fake_B = (fake_B_1 + fake_B_2)
            fake_B = generator(real_A)

            # fake_B = generator(fake_B)

            # pred_fake = discriminator(fake_B, real_A)
            pred_fake = discriminator(fake_B, real_B)
            
            # print(pred_fake.shape, valid.shape)
            loss_GAN = criterion_GAN(pred_fake, valid)
            # Pixel-wise loss
            loss_pixel = criterion_pixelwise(fake_B, real_B)

            # 特征提取器实现风格迁移损失

            # loss_style, loss_content = Feature(fake_B, real_A)
            # Style Loss
            # loss_style = feature_loss(fake_B, real_B, batch_size=opt.batch_size)


            tmp_real_B = norm_Zero2One(real_B, stratic=True)
            tmp_fake_B = norm_Zero2One(fake_B, stratic=True)

            # ssim_loss: torch.Tensor = piq.SSIMLoss(data_range=1.)(tmp_real_B, tmp_fake_B)
            # vif_loss: torch.Tensor = piq.VIFLoss(sigma_n_sq=2.0, data_range=1.)(tmp_real_B, tmp_fake_B)
            
            """
            图像边缘检测
            """
            # def edge_dect(im):
            #     # im = im.reshape((opt.batch_size, 1, im.shape[0], im.shape[1]))
            #     conv1 = nn.Conv2d(1, 1, 3, bias=False)
            #     sobel_kernel = torch.cuda.FloatTensor([[-1,	-1,	-1],	[-1,	8,	-1],	[-1,	-1,	-1]])
            #     sobel_kernel = sobel_kernel.reshape((1, 1, 3, 3))
            #     conv1.weight.data = sobel_kernel
            #     edge = conv1(Variable(im))
            #     return edge

            # def get_edge_loss(tmp_real_B, tmp_fake_B):
            #     edge_real_B = edge_dect(tmp_real_B)
            #     edge_fake_B = edge_dect(tmp_fake_B)
            #     loss_edge = criterion_pixelwise(edge_real_B, edge_fake_B)
            #     return loss_edge

            # edge_loss = get_edge_loss(tmp_real_B, tmp_fake_B)  
            

            # Total Loss
            # print("loss_GAN: {}, loss_pixel(with lambda): {}, loss_style: {}".format(loss_GAN.item(), lambda_pixel * loss_pixel.item(), loss_style.item()))
            loss_G = loss_GAN + lambda_pixel * loss_pixel  #+ loss_style + edge_loss # + vif_loss * 0.05  # + ssim_loss
        

            loss_G.backward()

            optimizer_G.step()



            # ---------------------
            #  训练判别器
            # ---------------------

            optimizer_D.zero_grad()

            # Real loss
            # pred_real = discriminator(real_B, real_A)
            pred_real = discriminator(fake_B.detach(), real_B)

            loss_real = criterion_GAN(pred_real, valid)

            # Fake loss
            pred_fake = discriminator(fake_B.detach(), real_A)
            # pred_fake = discriminator(fake_B.detach(), real_B)
            loss_fake = criterion_GAN(pred_fake, fake)

            # Total loss
            loss_D = loss_real + loss_fake

            loss_D.backward()
            # loss_D.backward(retain_graph=True)

            optimizer_D.step()

            # --------------
            #  Log Progress
            # --------------
            # tmp_real_B = norm_Zero2One(real_B, stratic=True)
            # tmp_fake_B = norm_Zero2One(fake_B, stratic=True)

            # mse, psnr = eval(real_B, fake_B)
            ssim = piq.ssim(tmp_fake_B, tmp_real_B, data_range=1.)
            ms_ssim = piq.multi_scale_ssim(tmp_fake_B, tmp_real_B, data_range=1.)
            psnr = piq.psnr(tmp_fake_B, tmp_real_B, data_range=1., reduction='none')
            vif = piq.vif_p(tmp_real_B, tmp_fake_B, data_range=1.)
            mse = (torch.abs(tmp_fake_B - tmp_real_B) ** 2)


            # ---------------------
            #  根据评估指标训练判别器
            # ---------------------



            # optimizer_G.zero_grad()
            # ssim_loss: torch.Tensor = piq.SSIMLoss(data_range=1.)(tmp_real_B, tmp_fake_B)
            # vif_loss: torch.Tensor = piq.VIFLoss(sigma_n_sq=2.0, data_range=1.)(tmp_real_B, tmp_fake_B)
            # print(ssim_loss, vif_loss)
            
            
            # total_eval_loss = ssim_loss + vif_loss
  
            # total_eval_loss.backward()

            # optimizer_G.step()




            all_loss_G += loss_G.item()
            all_loss_D += loss_D.item()
            all_loss_GAN += loss_GAN.item()
            all_loss_pixel += loss_pixel.item()
            # all_loss_vif += vif_loss.item()
            all_ms_ssim += ms_ssim.item()
            all_ssim += ssim.item()
            all_psnr += psnr.mean().item()
            all_vif += vif.item()
            all_mse += mse.mean().item()
            # all_edge_loss.append(edge_loss.item())

            # 输出训练进度
            # batches_done = epoch * len(dataloader) + i
            batches_done = epoch * opt.step + i
            # batches_left = opt.n_epochs * len(dataloader) - batches_done
            batches_left = opt.n_epochs * opt.step - batches_done

            time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
            prev_time = time.time()

            # If at sample interval save image
            if batches_done % opt.sample_interval == 0:
                sample_images(batches_done)

            # Print log
            sys.stdout.write(
                "\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, pixel: %f, adv: %f, edge_loss: no, vif_loss: no] [ssim: %f, ms_ssim: %f, psnr：%f, vif: %f, mse: %f]ETA: %s"
                % (
                    epoch,
                    opt.n_epochs,
                    i,
                    # len(dataloader),
                    opt.step, 
                    loss_D.item(),
                    loss_G.item(),
                    loss_pixel.item(),
                    loss_GAN.item(),
                    # edge_loss.item(), 
                    # vif_loss,
                    ssim.item(),
                    ms_ssim.item(), 
                    psnr.mean().item(),
                    vif.item(),
                    mse.mean().item(),
                    time_left,
                )
            )
        
        # TensorboardX  可视化训练过程
        train_num = i + 1
        # print("\n train_num: {}".format(train_num))
        # print('Loss_G: {}'.format(all_loss_G / train_num, epoch))
        # print('Loss_D: {}'.format(all_loss_D / train_num, epoch))
        # print('ssim: {}'.format(all_ssim / train_num, epoch))
        # print('mse: {}'.format(all_mse / train_num, epoch))
        # print('psnr: {}'.format(all_psnr / train_num, epoch))
        # print('vif: {}'.format(all_vif / train_num, epoch))


        writer.add_scalar('Loss_G', all_loss_G / train_num, epoch)
        writer.add_scalar('Loss_D', all_loss_D / train_num, epoch)
        # writer.add_scalar('ssim_loss', ssim_loss, epoch)
        writer.add_scalar('vif_loss', all_loss_vif / train_num, epoch)
        writer.add_scalar('ssim', all_ssim / train_num, epoch)
        writer.add_scalar('ms_ssim', all_ms_ssim / train_num, epoch)
        writer.add_scalar('mse', all_mse / train_num, epoch)
        writer.add_scalar('psnr', all_psnr / train_num, epoch)
        writer.add_scalar('vif', all_vif / train_num, epoch)
        writer.add_scalar('edge_loss', np.mean(all_edge_loss), epoch)
        if epoch % 10 == 0:
            with torch.no_grad():
                val(epoch)
        if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
            # Save model checkpoints
            torch.save(generator.state_dict(), "saved_models/%s/generator_%d.pth" % (opt.dataset_name, epoch))
            torch.save(discriminator.state_dict(), "saved_models/%s/discriminator_%d.pth" % (opt.dataset_name, epoch))
        
        scheduler_G.step()
        scheduler_D.step()
        writer.add_scalar('learning_rate', scheduler_G.get_last_lr(), epoch)

        adjust_learning_rate(epoch)
        torch.cuda.empty_cache()
    writer.close()